package ru.ifmo.network.skype.network;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import java.io.EOFException;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.security.*;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.X509EncodedKeySpec;

import static ru.ifmo.network.skype.network.NetworkUtils.*;

/**
 * @author avhaliullin
 */
public class RSAConnection implements SecureConnectionManager {


    private static final int KEY_SIZE = 1024;
    private static final PrivateKey privKey;
    private static final byte[] publicKey;
    private static final KeyFactory KEY_FACTORY;

    static {
        KeyPairGenerator kg = null;
        PublicKey pub;
        try {
            KEY_FACTORY = KeyFactory.getInstance("RSA");
            kg = KeyPairGenerator.getInstance("RSA");
            kg.initialize(KEY_SIZE);
            KeyPair pair = kg.generateKeyPair();
            pub = pair.getPublic();
            privKey = pair.getPrivate();
        } catch (NoSuchAlgorithmException e) {
            throw new AssertionError(e);
        }
        publicKey = pub.getEncoded();

    }

    @Override
    public Channel wrap(Channel channel) throws IOException {
        final OutputStream out = channel.getOutputStream();
        final InputStream in = channel.getInputStream();
        writeShort(out, (short) publicKey.length);
        out.write(publicKey);
        out.flush();

        int lenInt = readShort(in);
        byte[] otherPubKey = new byte[lenInt];
        readFull(in, otherPubKey);

        X509EncodedKeySpec pubKeySpec = new X509EncodedKeySpec(otherPubKey);
        try {
            PublicKey publicKey = KEY_FACTORY.generatePublic(pubKeySpec);
            Cipher decipher = Cipher.getInstance("RSA/ECB/NoPadding");
            decipher.init(Cipher.DECRYPT_MODE, privKey);
            Cipher cipher = Cipher.getInstance("RSA/ECB/NoPadding");
            cipher.init(Cipher.ENCRYPT_MODE, publicKey);
            return new RSAChannel(cipher, decipher, channel);
        } catch (InvalidKeySpecException | InvalidKeyException | NoSuchAlgorithmException | NoSuchPaddingException e) {
            throw new IOException(e);
        }
    }

    private class RSAChannel implements Channel {
        private final Cipher cipher;
        private final Cipher decipher;
        private final Channel channel;
        private final InputStream in;
        private final OutputStream out;
        private final InputStream secureIn = new InputStream() {

            @Override
            public int read() throws IOException {
                if (inBufDecrypted == null || inOffset >= msgSize) {
                    doUpdate();
                    if (inBufDecrypted == null) {
                        return -1;
                    }
                }
                int res = inBufDecrypted[inOffset++];
                return res >= 0 ? res : (res) + 256;
            }
        };

        private final OutputStream secureOut = new OutputStream() {

            @Override
            public void close() throws IOException {
                doFlush();
            }

            @Override
            public void flush() {
                try {
                    doFlush();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }

            @Override
            public void write(int b) throws IOException {
                outBuf[outOffset++] = (byte) (b < 0x7F ? b : b - 256);
                if (outOffset >= outBuf.length) {
                    doFlush();
                }
            }
        };

        public RSAChannel(Cipher cipher, Cipher decipher, Channel channel) throws IOException {
            this.cipher = cipher;
            this.decipher = decipher;
            this.in = channel.getInputStream();
            this.out = channel.getOutputStream();
            this.channel = channel;
        }

        private void doFlush() throws IOException {
            try {
                writeShort((short) outOffset, outBuf, 0);
                byte[] res = cipher.doFinal(outBuf);
                out.write(res);
                out.flush();
                outOffset = 2;
            } catch (IllegalBlockSizeException | BadPaddingException e) {
                throw new AssertionError(e);
            }
        }

        private final byte[] outBuf = new byte[KEY_SIZE / 8];
        private int inOffset = 0;
        private int outOffset = 2;
        private int msgSize = 0;

        private byte[] inBuf = new byte[KEY_SIZE / 8];
        private byte[] inBufDecrypted;

        private void doUpdate() throws IOException {
            inBuf = new byte[KEY_SIZE / 8];
            try {
                readFull(in, inBuf);
                try {
                    inBufDecrypted = decipher.doFinal(inBuf);
                    inOffset = 2;
                    msgSize = readShort(inBufDecrypted, 0);
                } catch (IllegalBlockSizeException | BadPaddingException e) {
                    throw new AssertionError(e);
                }
            } catch (EOFException e) {
                inBufDecrypted = null;
            }
        }

        @Override
        public InputStream getInputStream() throws IOException {
            return secureIn;
        }

        @Override
        public OutputStream getOutputStream() throws IOException {
            return secureOut;
        }

        @Override
        public void close() {
            try {
                doFlush();
            } catch (IOException e) {
            }
            channel.close();
        }
    }
}
